{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ifmV02zKlsCs",
        "outputId": "57c94cfd-234d-4c47-8cfb-12730e0ae793"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m493.7/493.7 kB\u001b[0m \u001b[31m7.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m7.9/7.9 MB\u001b[0m \u001b[31m37.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m115.3/115.3 kB\u001b[0m \u001b[31m13.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m134.8/134.8 kB\u001b[0m \u001b[31m17.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m311.2/311.2 kB\u001b[0m \u001b[31m36.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m3.8/3.8 MB\u001b[0m \u001b[31m78.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m74.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u2501\u001b[0m \u001b[32m295.0/295.0 kB\u001b[0m \u001b[31m33.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "!pip install -qU \\\n",
        "  datasets==2.14.6 \\\n",
        "  transformers==4.35.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q4vTJ-pFmWl5"
      },
      "source": [
        "## Dataset Download"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GFaaDw5VmZEk"
      },
      "source": [
        "We're going to test with a more real world use-case, with messy, imperfect data. We will use the [`jamescalam/ai-arxiv-chunked`](https://huggingface.co/datasets/jamescalam/ai-arxiv-chunked) dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 234,
          "referenced_widgets": [
            "9797e8b888524e859e0186488e768b97",
            "90eae7e29b2f4b42affbc23f7a45d15d",
            "0596666289d14478a5ffe889ce518dbc",
            "e9014129c4f048f99186087c29527a95",
            "53da138dcde44acf88741ff8b7841017",
            "51b19cf22e144030a72dbd0288d1e8e8",
            "ef89bbd0b01c4e3ca44d1dcbb5e88a26",
            "c6dbf744efd54029a510824bd0fab5b6",
            "bacb32bc40044d9b88b184f8f5d5523a",
            "c2ff3d2303de44d983b7402c484bb06c",
            "af87203a68ab4fd08a1e079a74139be3",
            "94943162bb9d46199dc3e7d40d012b1b",
            "178982a6664546fcb4deccd9f9f93309",
            "f8bcf759f42246f992430d57d7cc53e6",
            "c9ecf805aaa04166a61df32280b3e771",
            "d6f09ae72c014e1cbb9dbe71bf8640ea",
            "f3526a6c70dd4aada5fbcbc1debb72e5",
            "2858afd141e944e3bcc53ca65e8dc296",
            "2e92df0f3506404ab428e2e22e822ac7",
            "527364a2b7f341d7bf58812a217992c1",
            "340886a1803f4034b77c194ed9b8e06d",
            "556f2aaad36545aeb1c1a2048dcc77c5",
            "751be32f96dd4e8abd0c2e8daa496319",
            "fb3f34060f2d4d29b1a7898e2e7f6a86",
            "cc1cae1de99248c982eb287d277aa89b",
            "65e314995f50482baad3579b54b15d40",
            "4d06051ffe584f8b9e087dc529ec0f05",
            "f87edae5b3b14b0eb7f56ece0532972f",
            "dd7a6493b8454422ba3ff66a4706b37d",
            "9f82f72395f54c22bb0de3d39e0ab706",
            "23937cf56f1e4c269523bf7090089f62",
            "49feccd41ddb43eb9a3da710cae5649c",
            "e923ae7b8f8347838ae2f95dc28f9974",
            "39c8e889b8f34593b87548edbf8c9d23",
            "4cd201469ca14904a54115c3878dba8d",
            "18d03556254140d8a50c8c3c4af5f8a0",
            "57a67e93ecaa41e88d574a870b352e7b",
            "3d757b79b68743129f59c8117954a597",
            "864f76dc043e4cf2b9c537223b40e81b",
            "1fc2cba2e61142628e0f60cd4e6ffc13",
            "2068a8781b534ba6b7eeefc994eff119",
            "3a2c3b2dba3a4af4b36f27742d4a68ee",
            "2f61bcdd61d04b94bf610f7497752037",
            "7baf722b727f466ca7dbb0a69d3432fb"
          ]
        },
        "id": "4-FqcdKHmVpa",
        "outputId": "7350e758-639e-4988-9d7b-4e36d058b3ba"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading data files:   0%|          | 0/1 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "9797e8b888524e859e0186488e768b97"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading data:   0%|          | 0.00/153M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "94943162bb9d46199dc3e7d40d012b1b"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Extracting data files:   0%|          | 0/1 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "751be32f96dd4e8abd0c2e8daa496319"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Generating train split: 0 examples [00:00, ? examples/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "39c8e889b8f34593b87548edbf8c9d23"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Dataset({\n",
              "    features: ['doi', 'chunk-id', 'chunk', 'id', 'title', 'summary', 'source', 'authors', 'categories', 'comment', 'journal_ref', 'primary_category', 'published', 'updated', 'references'],\n",
              "    num_rows: 41584\n",
              "})"
            ]
          },
          "metadata": {},
          "execution_count": 2
        }
      ],
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "data = load_dataset(\"jamescalam/ai-arxiv-chunked\", split=\"train\")\n",
        "data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gp5a_bInyfdX"
      },
      "source": [
        "First we define our embedding function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "oG6zd1dLw54w",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 227,
          "referenced_widgets": [
            "f7a2cd801be3427d8e92bd575710ae89",
            "dac571167e01494481ba2c8971ec89e1",
            "262cac9f0d60471eafc0cd7688f3f8e6",
            "2b18265f5d034494a1f1bc66dad77328",
            "91d32a1a5d5b4a48a1cc8fd4ca0f961f",
            "a810c8ccbedf425f9f3002740e9c8cf8",
            "ec07e6bf8af14d969a855e61eb9bb597",
            "553f848da3614b8d86dcc9ed13a1cc5e",
            "76818e4ee4f04049896df1542e5d0ddd",
            "a6e2f44107e8484b90b8e183e78f5552",
            "19118c4a533e45d89d5b4068e39f1689",
            "5bb3df63978842d096e79e78919dd6a6",
            "5c6c8fed929a4e70943d10221fa036fe",
            "96c927faa2d44760b13e20e8a03301d4",
            "378bb53aaccd490296cd741d013e5a42",
            "344d15a908164868808ab59c441b728f",
            "84ac8e5cbcdc4206ad55b4020e9dd5c6",
            "4c6d8f438a114858b07e7c802ce7bea0",
            "99327b3dc0fa4eca963860f928801183",
            "f70f383b6d9c41ecad01012b279cf05a",
            "0d15a123f80246af9de060c9c074e932",
            "3845fd0c4e024348bf5218c4878ef0fe",
            "7812980cad8c45fda6d8308851247e39",
            "c8ba1592cbed45498dcdaaea850c6a02",
            "5b72a3db5d824ee19cc4e92be0ae1e50",
            "43693813d4074a529edc8a300a5f139d",
            "5ba3b8f0924a487086b64e2cafaddc0b",
            "8d1692efe0d64c3fb8cd6d4ab0dec986",
            "c15f8724fd6f417589995c4409e9314d",
            "84c6322b030442129d6d7820f9bae428",
            "941f30092dee490ca5ac061611862471",
            "cbd5d36be3d74916841f4d21a5495263",
            "c45dae576fc94786b09e3e2fd2aefa68",
            "47b86f27c0c84fdaa6edfb54e35eb9fc",
            "36a26d8b0aa344cda83387d6853826a0",
            "af31532c4a924cfb953b7b851b9482ef",
            "31fa48d3976449d6896e34e662585e46",
            "e876d7a30aa8433b96a2f4284b57985e",
            "c61bf77a1ebf4918a2ba90b3c763935a",
            "12d30c766a194f49ad3959f789e97077",
            "6830c85d0ebe49c0862e5358e1eb7838",
            "225c776ffd3a4eb18bda226eae27b00d",
            "ea033cf4cfeb4828b14064cb0e9d2c6a",
            "437b9ebcd86b4ae2996bbc62a7cdbda8",
            "1d5b40f1b76244d6826714f76971a0e5",
            "48c55f5652e54dc986a6e7728a4f03d8",
            "d88fda6f75344317968b353ac9baaef1",
            "66ddca296890418ca8caf360bdc7012d",
            "905232c025b34af5ad87b70446f77d82",
            "61db3be0156247ec903866be621f8d74",
            "774b01cbbf4243088e49158c7f92ab92",
            "b99995169c9545e0b4598dc3cdef47c6",
            "811697672c30463aa2206737e4285ec8",
            "b93be5b5d8a04c05859c87048826fd34",
            "810f978cebf94c6eafb49d0ec4811545",
            "0e25afc0dded4e449682190deefdde9d",
            "042b7f9e0bec49f8b455b567ee420f55",
            "abb4456b7b0544a991c63096bd181f91",
            "fdeb2bf7a6a24eac9d57385f459e0554",
            "1e0df624a27043b29155a3eef7116ff6",
            "35bff7142e4c48cf84e9a204193ad9f5",
            "81e030ce95344a13b45dd9e16bd990e2",
            "e9d3c9b4d1254864b3c0111aedaf1362",
            "b6a2acb15fb14f7e968cac86844ee904",
            "b490d64e539048d79dfbc066ecdfb2bf",
            "1a67c85fa7a04361941c0b84058568ce"
          ]
        },
        "outputId": "a529ae4f-a4c4-4a79-e935-f27ef457f7ff"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Using cuda\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading (\u2026)okenizer_config.json:   0%|          | 0.00/366 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "f7a2cd801be3427d8e92bd575710ae89"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading (\u2026)solve/main/vocab.txt:   0%|          | 0.00/232k [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "5bb3df63978842d096e79e78919dd6a6"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading (\u2026)/main/tokenizer.json:   0%|          | 0.00/711k [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "7812980cad8c45fda6d8308851247e39"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading (\u2026)cial_tokens_map.json:   0%|          | 0.00/125 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "47b86f27c0c84fdaa6edfb54e35eb9fc"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading (\u2026)lve/main/config.json:   0%|          | 0.00/696 [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "1d5b40f1b76244d6826714f76971a0e5"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Downloading pytorch_model.bin:   0%|          | 0.00/219M [00:00<?, ?B/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "0e25afc0dded4e449682190deefdde9d"
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "import torch\n",
        "from torch.nn.functional import normalize\n",
        "from transformers import AutoModel, AutoTokenizer\n",
        "\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "print(f\"Using {device}\")\n",
        "\n",
        "model_id = \"infgrad/stella-base-en-v2\"\n",
        "\n",
        "# initialize tokenizer and model\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
        "model = AutoModel.from_pretrained(model_id).to(device)\n",
        "model.eval()\n",
        "\n",
        "def embed(docs: list[str]) -> list[list[float]]:\n",
        "    # tokenize\n",
        "    tokens = tokenizer(\n",
        "        docs, padding=True, max_length=512, truncation=True, return_tensors=\"pt\"\n",
        "    ).to(device)\n",
        "    with torch.no_grad():\n",
        "        # process with model for token-level embeddings\n",
        "        out = model(**tokens)\n",
        "        # mask padding tokens\n",
        "        last_hidden = out.last_hidden_state.masked_fill(\n",
        "            ~tokens[\"attention_mask\"][..., None].bool(), 0.0\n",
        "        )\n",
        "        # create mean pooled embeddings\n",
        "        doc_embeds = last_hidden.sum(dim=1) / \\\n",
        "            tokens[\"attention_mask\"].sum(dim=1)[..., None]\n",
        "    return doc_embeds.cpu().numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1nvrNQSGXvEC"
      },
      "source": [
        "Use this to build a Numpy array of cohere embedding vectors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "EdyWVR17zX7I",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "52844fc77f4b426d8087f8ca031a0687",
            "652587b101f743efa6fd1a71035e1baf",
            "047c8e9315d64e83af47d67da7f9631b",
            "851da7f83a104c2d810d3f73a038f757",
            "0f5b416760ad4245aa7e63a790cf26d2",
            "b127dd6a4c3243a081515a7113bbec08",
            "8f0afd50a72e436c99ff778ad3a10a9f",
            "b0bdcc0c58f145b6885b48157dcc9340",
            "93ae9c88ac0849cbbe0dd6150a670df1",
            "cb19c82b6e404086adc5d1b133d3f72a",
            "30bd1f9db48d476e8f7a301c0cb64c5c"
          ]
        },
        "outputId": "0954874d-fe38-42ed-c9db-486de0206e3d"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/325 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "52844fc77f4b426d8087f8ca031a0687"
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "from tqdm.auto import tqdm\n",
        "import numpy as np\n",
        "\n",
        "chunks = data[\"chunk\"]\n",
        "batch_size = 128\n",
        "\n",
        "for i in tqdm(range(0, len(chunks), batch_size)):\n",
        "    i_end = min(len(chunks), i+batch_size)\n",
        "    chunk_batch = chunks[i:i_end]\n",
        "    # embed current batch\n",
        "    embed_batch = embed(chunk_batch)\n",
        "    # add to existing np array if exists (otherwise create)\n",
        "    if i == 0:\n",
        "        arr = embed_batch.copy()\n",
        "    else:\n",
        "        arr = np.concatenate([arr, embed_batch.copy()])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bl9g3ePt029u"
      },
      "source": [
        "Now we need to create the query mechanism, this is simply a cosine similarity calculation between a query vector and our `arr` vectors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "MR7WyDiatlsX"
      },
      "outputs": [],
      "source": [
        "from numpy.linalg import norm\n",
        "\n",
        "# convert chunks list to array for easy indexing\n",
        "chunk_arr = np.array(chunks)\n",
        "\n",
        "def query(text: str, top_k: int=3) -> list[str]:\n",
        "    # create query embedding\n",
        "    xq = embed([text])[0]\n",
        "    # calculate cosine similarities\n",
        "    sim = np.dot(arr, xq.T) / (norm(arr, axis=1)*norm(xq.T))\n",
        "    # get indices of top_k records\n",
        "    idx = np.argpartition(sim, -top_k)[-top_k:]\n",
        "    docs = chunk_arr[idx]\n",
        "    for d in docs.tolist():\n",
        "        print(d)\n",
        "        print(\"----------\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "u2NjYxsn7J5f",
        "outputId": "5528c0c3-09e8-4814-d9a7-6dff5f62d251"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "chat with itself. For future work, we would like\n",
            "to explore introducing reinforcement learning to\n",
            "further improve the performance of our models.\n",
            "Limitations\n",
            "Foundation Model Similar to other language\n",
            "models, Baize may suffer from hallucination, toxicity and stereotypes. Particularly, Baize inherits the\n",
            "out-of-date knowledge from LLaMA. Due to the\n",
            "fact that at least 82% of LLaMA\u2019s pretraining data\n",
            "is from before 2020, Baize may provide outdated\n",
            "answers to certain questions, such as \"who is the\n",
            "current president of the United States?\" Additionally, LLaMA only supports 20 languages and has a\n",
            "very limited corpus for non-English languages.\n",
            "Evaluation In this paper, we automatically evaluating the models with GPT-4 (OpenAI, 2023b).However, we found that it has a strong preference\n",
            "for longer responses and a positional bias. We believe human evaluation can be more rigorous and reliable despite being expensive and time-consuming\n",
            "while automatic evaluation remains an open research question.\n",
            "License and Legality Following Stanford Alpaca (Taori et al., 2023), we have decided that the\n",
            "released weights of Baize are licensed for research\n",
            "use only. Using the weights of Baize with LLaMA\u2019s\n",
            "original weights is subject to Meta\u2019s LLaMA License Agreement. It is the responsibility of the\n",
            "----------\n",
            "PaLM8B 25.6 23.8 24.1 27.8 25.4\n",
            "62B 59.5 41.9 62.7 55.8 53.7\n",
            "540B 77.0 55.6 81.0 69.6 69.3\n",
            "LLaMA7B 34.0 30.5 38.3 38.1 35.1\n",
            "13B 45.0 35.8 53.8 53.3 46.9\n",
            "33B 55.8 46.0 66.7 63.4 57.8\n",
            "65B 61.8 51.7 72.9 67.4 63.4\n",
            "Table 9: Massive Multitask Language Understanding (MMLU). Five-shot accuracy.\n",
            "that may indicate that this benchmark is not\n",
            "reliable. On WinoGrande, the performance does\n",
            "not correlate as well with training perplexity:\n",
            "the LLaMA-33B and LLaMA-65B have similar\n",
            "performance during the training.\n",
            "4 Instruction Finetuning\n",
            "In this section, we show that brie\ufb02y \ufb01netuning on\n",
            "instructions data rapidly leads to improvements\n",
            "on MMLU. Although the non-\ufb01netuned version\n",
            "of LLaMA-65B is already able to follow basic instructions, we observe that a very small amount of\n",
            "\ufb01netuning improves the performance on MMLU,\n",
            "----------\n",
            "\u0003Equal contribution. Correspondence: {htouvron,\n",
            "thibautlav,gizacard,egrave,glample}@meta.com\n",
            "1https://github.com/facebookresearch/llamaperformance, a smaller one trained longer will\n",
            "ultimately be cheaper at inference. For instance,\n",
            "although Hoffmann et al. (2022) recommends\n",
            "training a 10B model on 200B tokens, we \ufb01nd\n",
            "that the performance of a 7B model continues to\n",
            "improve even after 1T tokens.\n",
            "The focus of this work is to train a series of\n",
            "language models that achieve the best possible performance at various inference budgets, by training\n",
            "on more tokens than what is typically used. The\n",
            "resulting models, called LLaMA , ranges from 7B\n",
            "to 65B parameters with competitive performance\n",
            "compared to the best existing LLMs. For instance,\n",
            "LLaMA-13B outperforms GPT-3 on most benchmarks, despite being 10 \u0002smaller. We believe that\n",
            "this model will help democratize the access and\n",
            "study of LLMs, since it can be run on a single GPU.\n",
            "At the higher-end of the scale, our 65B-parameter\n",
            "model is also competitive with the best large language models such as Chinchilla or PaLM-540B.\n",
            "Unlike Chinchilla, PaLM, or GPT-3, we only\n",
            "----------\n"
          ]
        }
      ],
      "source": [
        "query(\"why should I use llama 2?\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "ZSIBTMUy7qc0",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "96af6223-5ea1-45a3-965b-1837a10ef907"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "for red teaming (\u00a73). Throughout the design of our experiments, we arrived at many junctures in which\n",
            "we were unsure about how to proceed, even after a literature review on red teaming AI systems (\u00a72). As\n",
            "such, we conducted informational interviews with experts in the \ufb01eld of Trust & Safety and incorporated\n",
            "their suggested best practices (\u00a7A.2) into the design of our experiments in order to ensure the well-being of\n",
            "the red team. In general, we found that red team members enjoyed participating in our experiments and felt\n",
            "motivated by a mission to make AI systems less harmful (\u00a7A.2). Nevertheless, our work suffers from some\n",
            "limitations, which we discuss in \u00a75.1. Based on our experiences, we propose some policy interventions for\n",
            "how we can work together as a community to develop shared norms, practices, and technical standards for\n",
            "how to red team language models (\u00a75.2).\n",
            "2 Related Work\n",
            "We use the same models that we developed in our previous work where we train a general language assistant\n",
            "to be helpful, honest, and harmless [2, 4]. However, here we run additional experiments in order to determine\n",
            "the in\ufb02uence of model size on susceptibility to red team attacks (Figure 1) and analyze the content of the\n",
            "----------\n",
            "AI research community could take to build consensus around how to red team andhow to release \ufb01ndings\n",
            "from red teaming .\n",
            "For how to red team , we have detailed our initial approach. However, we conducted this effort in isolation, and we would have bene\ufb01ted from participating in a community-based effort to address certain open\n",
            "questions:\n",
            "\u2022 Who should red team and why?\n",
            "\u2022 What protections should we put in place to ensure the safety of the red team?\n",
            "\u2022 What instructions and information about the models should we provide to the red team?\n",
            "\u2022 How should we annotate and analyze the data we collect?\n",
            "\u2022 What constitutes a successful red team attempt?\n",
            "We can make progress towards answering these questions by convening a multidisciplinary community to\n",
            "share different approaches to internal red teaming and drive toward consensus.\n",
            "The research community lacks shared norms and best practices for how to release \ufb01ndings from red teaming. As a result, we made our decision to release the data largely on our own and likely missed critical\n",
            "perspectives from experts, other disciplines, and members of the public.14The decision for how to appropriately release \ufb01ndings will ultimately require a subjective judgment call. For our purposes, we reviewed a\n",
            "sample of our red team dataset and evaluated the pros and cons of a public release (See \u00a7A.5). Among them\n",
            "----------\n",
            "respectively. We could address this problem by asking third party organizations that specialize in certain\n",
            "domains to red team our systems. Additionally, we could give crowdworkers a way to indicate if and how a\n",
            "particular red team attack requires domain expertise in order to evaluate how successful the attack was.\n",
            "As expected, our data are incomplete\u2014because LMs are general purpose and open-ended, the space of possible harms is unknown and unbounded [22]. For example, the models we red teamed have been trained partly\n",
            "on Python code; however, we observed no attacks related to code generation ability (e.g., \u201cwrite a Python\n",
            "program that implements a DDOS attack\u201d). It is possible that sharing our red team interface with more domain experts could have resulted in such attacks. We could have also noted in the instructions to the interface\n",
            "that such attacks would be viable, but we erred on the side of being less prescriptive about how to red team\n",
            "in order to encourage creativity. It is unclear how to strike the right balance.\n",
            "We also know our data are incomplete because we informally red teamed our models internally and found\n",
            "successful attack types not present in the dataset we release. For example, we uncovered a class of attacks\n",
            "that we call \u201croleplay attacks\u201d on the RLHF model. In a roleplay attack we exploit the helpfulness of the\n",
            "----------\n"
          ]
        }
      ],
      "source": [
        "query(\"can you tell me about red teaming for llama 2?\")"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "query(\"what is the best llm?\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2VUqwF-com3z",
        "outputId": "72974974-19aa-42a6-8a4b-665a7d1a2d6e"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "as well as the amount of data required to outperform LLMs. We surpass the performance of 540B\n",
            "parameter LLMs using a 770M T5 model; this\n",
            "smaller model only uses 80% of a labeled dataset\n",
            "that would otherwise be required if using an exist-ing \ufb01netuning method. When only unlabeled data\n",
            "is present, our small models still perform on par or\n",
            "better than LLMs. We outperform 540B PaLM\u2019s\n",
            "performance with only a 11B T5 model. We further\n",
            "show that when a smaller model performs worse\n",
            "than an LLM, Distilling step-by-step can more ef\ufb01ciently leverage additional unlabeled data to match\n",
            "the LLM performance compared to the standard\n",
            "distillation approach.\n",
            "2 Related work\n",
            "Our work distills task-speci\ufb01c knowledge of LLMs\n",
            "into smaller specialist models by leveraging the\n",
            "emergent reasoning capabilities of today\u2019s LLMs.\n",
            "We draw on knowledge distillation research and\n",
            "methods that learn from both human-generated rationales and LLM-generated rationales.\n",
            "Knowledge distillation from large models.\n",
            "Knowledge distillation has been successfully used\n",
            "to transfer knowledge from larger, more competent teacher models into smaller student models\n",
            "affordable for practical applications (Bucilu \u02c7a et al.,\n",
            "----------\n",
            "LLM outperforms the 280B parameter LLM (57.31% vs 51.45%), while still performing signi\ufb01cantly\n",
            "worse than the 7B SI model ( \ud835\udc5d=0\u0093012). Note that the latter is evaluated in the harder generative\n",
            "setting. Per task breakdown shown in Fig. 4b demonstrates that the SI framework solves the bAbI\n",
            "15 deduction task, the only model to achieve 100% accuracy (signi\ufb01cant di\ufb00erence from the other\n",
            "models,\ud835\udc5d \u009d0\u009301). Furthermore, it does so having seen only \ufb01ve examples in the prompt. The 7B SI\n",
            "modelalsosigni\ufb01cantlyoutperformsallothermodelsonProofWriterDepth0( \ud835\udc5d \u009d0\u009301),ProofWriter\n",
            "Depth 1 (\ud835\udc5d=0\u0093034).\n",
            "As well as improving upon most baselines quantitatively, the SI framework also has additional\n",
            "qualitative bene\ufb01ts: 1) it produces a causal, human interpretable reasoning trace that shows how the\n",
            "model reached its answer and 2) it is able to recover from errors. We will now discuss each of these\n",
            "3Thiscouldsuggest that the 280B LLM has stronger priors, than the 7B LLM, which it favours over logical reasoning. For\n",
            "----------\n",
            "for LLMs has been concentrated on prompting and in-context learning due to their tremendous\n",
            "sizes (Brown et al., 2020; Liu et al., 2021a). Following peer LLMs, in this work our evaluation is\n",
            "also under such setting. Nevertheless, some recent attempts has been on parameter-ef\ufb01cient learning\n",
            "on language models (Houlsby et al., 2019) and prompt tuning (i.e., P-tuning, Li & Liang (2021);\n",
            "Liu et al. (2021c); Lester et al. (2021); Liu et al. (2022)). In this work, we do not focus on them and\n",
            "will leave their testing on GLM-130B in future study.\n",
            "Inference. Most public-accessible LLMs nowadays are providing their services via limited APIs.In\n",
            "this work, an important part of our endeavor has been on LLMs\u2019 ef\ufb01cient and fast inference. Related\n",
            "work may include distillation (Sanh et al., 2019; Jiao et al., 2020; Wang et al., 2020), quantization (Zafrir et al., 2019; Shen et al., 2020; Tao et al., 2022), and pruning (Michel et al., 2019; Fan\n",
            "et al., 2019). Very recent work (Dettmers et al., 2022) shows that LLMs such as OPT-175B and\n",
            "----------\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "query(\"what is the difference between gpt-4 and llama 2?\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sF8xvtyaoom4",
        "outputId": "14d60aac-130e-44ba-af1e-d58908b63638"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "to GPT-3 corresponds to the Stanford Alpaca model. From Figure 3(a), we observe that ( i) For the\n",
            "\u201cHelpfulness\u201d criterion, GPT-4 is the clear winner with 54.12% of the votes. GPT-3 only wins 19.74%\n",
            "of the time. ( ii) For the \u201cHonesty\u201d and \u201cHarmlessness\u201d criteria, the largest portion of votes goes\n",
            "to the tie category, which is substantially higher than the winning categories but GPT-3 (Alpaca) is\n",
            "slightly superior.\n",
            "Second, we compare GPT-4-instruction-tuned LLaMA models against the teacher model GPT-4 in\n",
            "Figure 3(b). The observations are quite consistent over the three criteria: GPT-4-instruction-tuned\n",
            "LLaMA performs similarly to the original GPT-4. We conclude that learning from GPT-4 generated\n",
            "5\n",
            "60% 70% 80% 90% 100%12345BRanking Group 94% 624 : 66792% 614 : 67091% 623 : 68289% 597 : 66989% 605 : 67891% 609 : 666\n",
            "----------\n",
            "tasks.\n",
            "This represents work in progress, and several directions can be explored: (i)Data and model scale .\n",
            "The GPT-4 data size is 52K and the base LLaMA model size is 7B. Vicuna collects around 700K\n",
            "conversion turns (approximated from the multi-turn ShareGPT data), and uses the 13B LLaMA\n",
            "model. Therefore, it would be promising to continue collecting more GPT-4 instruction-following\n",
            "data, combine with ShareGPT data, and train larger LLaMA models for higher performance. (ii)\n",
            "RLHF . The reward model is only used in the decoding stage, which suggests that comparison data is\n",
            "promising to provide useful feedback for LLM training. It is natural to continue to train LLMs with\n",
            "reward models, for example for reinforcement learning using machine-generated feedback.\n",
            "8\n",
            "ACKNOWLEDGMENTS\n",
            "We thank Guoyin Wang, Haotian Liu and Hao Cheng for valuable discussions and insightful experience sharing on instruction-tuning language models. We thank the LLaMA team for giving us access\n",
            "to their models.\n",
            "REFERENCES\n",
            "Amanda Askell, Yuntao Bai, Anna Chen, Dawn Drain, Deep Ganguli, Tom Henighan, Andy Jones,\n",
            "Nicholas Joseph, Ben Mann, Nova DasSarma, et al. A general language assistant as a laboratory\n",
            "----------\n",
            "-0.043\n",
            "-0.009+0.0132-0.004 +0.0562\n",
            "+0.0387-0.012\n",
            "-0.076Alpaca: 0.39 LLaMA-GPT4: 0.34 GPT4: 0.37Figure 6: ROUGE-L on unnatural instructions evaluated with 9K samples. The instructions are\n",
            "grouped into four subsets based on the ground-truth response length. The mean values are reported in\n",
            "the legend. The difference with GPT-4 is reported on the bar per group. LLaMA-GPT4 is a closer\n",
            "proxy to GPT-4 than Alpaca.\n",
            "closely follow the behavior of GPT-4. When the sequence length is short, both LLaMA-GPT4 and\n",
            "GPT-4 can generate responses that contains the simple ground truth answers, but add extra words to\n",
            "make the response more chat-like, which probably leads to lower ROUGE-L scores.\n",
            "5 R ELATED WORK\n",
            "Instruction Tuning. Instruction tuning of LLMs is an increasingly popular research direction in\n",
            "NLP (Zhong et al., 2021; Ouyang et al., 2022; Wei et al., 2021). Existing works aim to improve\n",
            "the quality and scale of three factors in the development pipeline, including instruction-following\n",
            "----------\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "---"
      ],
      "metadata": {
        "id": "d2UP-QJzo610"
      }
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "machine_shape": "hm",
      "provenance": [],
      "gpuType": "V100"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}